% 定性展示氢原子的各个电子轨道
% 已知bug:在原点附近的Psi计算可能有错误.
% Gitee Repo

clc
clear

a = 0.2;

n = 4;  %主量子数 n=1,2,3,..., KLMN
l = 3;  %角量子数 l=0,1,2,...,n-1, spdf
m = 0;  %磁量子数 m =0,1,2,...,l

dx = 0.05;
L = 5;

function R = compute_R(r, a, n, l)
  kappa = 1/(a*n);
  rho = kappa * r;

  R = 1;
  c = zeros(1,n);

  for j = 1:n-l-1;
    C = 2*(j+l-n)/(j*(j+2*l+1));
    if j == 1
      c(j) = C*1;
    else
      c(j) = C*c(j-1);
    end
    R += c(j)*(rho.^j);
  end
  R = (1./r).*rho.^(l+1).*(e.^(-rho)).*R;
end

function Y=compute_Y(x,l,m)
    [sx sy sx] = size(x);
    _Y = legendre(l,x);
    Y = zeros(sx,sy,sx);
    Y(:,:,:) = _Y(m+1,:,:,:);
end

if n <= 0
    error('invalid n')
    return
end
if l < 0 || l >= n
    error('invalid l')
    return
end
if m < 0 || m > l
    error('invalid m')
    return
end

[x y z] = meshgrid(-L:dx:L);
R = sqrt(x.^2+y.^2+z.^2);
costheta = z./R;

Rnl = compute_R(R, a, n, l); %径向方程
Ylm = compute_Y(costheta,l,m); %角方程
Psi = Rnl.*Ylm;
Psi = Psi.^2;

Psi(isnan(Psi) | isinf(Psi)) = max(Psi(:));

Psi=Psi/(sum(Psi(:))*(dx)^3);
Psi(Psi>0.05)=0.05; %优化显示，避免过大的峰

figure();
hold on
axis equal
h = slice(x,y,z,Psi,[],0,[]);
%caxis([0 1])
view(0,0)
set(h,'edgecolor','none')

% 导出数据以便使用其他软件绘图
save('atomic_orbs.mat','x','y','z','Psi','-v7')
